import akshare as ak
import pandas as pd
import time
import os
import traceback

# 输入输出路径
input_file = "../basic_info/code_name.csv"
output_file = "market_data_20250805.csv"
failed_log = "failed_codes.txt"
date = "20250805"

# 输出字段列（增加“股票简称”）
output_columns = [
    "日期", "股票代码", "股票简称", "开盘", "收盘", "最高", "最低",
    "成交量", "成交额", "振幅", "涨跌幅", "涨跌额", "换手率"
]

# 获取已写入的股票代码集合
if os.path.exists(output_file):
    written_df = pd.read_csv(output_file, dtype={"股票代码": str})
    written_codes = set(written_df["股票代码"].unique())
    print(f"已写入股票数量: {len(written_codes)}")
else:
    pd.DataFrame(columns=output_columns).to_csv(output_file, index=False, encoding="utf-8-sig")
    print(f"已创建输出文件并写入表头：{output_file}")
    written_codes = set()

# 读取股票代码列表（保留简称）
df_code = pd.read_csv(input_file, dtype={"股票代码": str})

max_retries = 10

# 遍历每只股票
for _, row in df_code.iterrows():
    code = row["股票代码"]
    name = row["股票简称"]

    if code in written_codes:
        print(f"{code} 已存在，跳过")
        continue

    for attempt in range(1, max_retries + 1):
        try:
            df = ak.stock_zh_a_hist(symbol=code, period="daily", start_date=date, end_date=date, adjust="")
            
            if df.empty:
                print(f"{code} 无数据，跳过")
                break

            data = df.iloc[0]
            new_row = {
                "日期": data["日期"],
                "股票代码": code,
                "股票简称": name,
                "开盘": data["开盘"],
                "收盘": data["收盘"],
                "最高": data["最高"],
                "最低": data["最低"],
                "成交量": data["成交量"],
                "成交额": data["成交额"],
                "振幅": data["振幅"],
                "涨跌幅": data["涨跌幅"],
                "涨跌额": data["涨跌额"],
                "换手率": data["换手率"]
            }

            pd.DataFrame([new_row]).to_csv(output_file, mode="a", header=False, index=False, encoding="utf-8-sig")
            print(f"{code} 写入成功")
            time.sleep(1)
            break
        except Exception as e:
            print(f"{code} 第 {attempt} 次尝试失败: {e}")
            if attempt < max_retries:
                time.sleep(1)
            else:
                print(f"{code} 最终获取失败，写入失败日志")
                with open(failed_log, "a") as f:
                    f.write(f"{code}\n")
                traceback.print_exc()
